import pytest  # used for our unit tests
from workload import AlexNet

# unit tests

# --- BASIC TEST CASES ---

def test_extract_features_single_image_exact_size():
    # Test with a single image whose flattened size matches features_size
    model = AlexNet()
    img = [list(range(1000))]  # single image, already flat
    codeflash_output = model._extract_features(img); features = codeflash_output # 982ns -> 411ns (139% faster)

def test_extract_features_multiple_images_varied_shapes():
    # Test with multiple images, each a list of lists (simulate 2D images)
    model = AlexNet()
    img1 = [[i for i in range(36)] for _ in range(256)]  # 256x36 = 9216
    img2 = [[1]*36 for _ in range(256)]
    codeflash_output = model._extract_features([img1, img2]); features = codeflash_output # 911ns -> 431ns (111% faster)

def test_extract_features_shorter_than_features_size():
    # Test with an image whose flattened size is less than features_size (should pad with zeros)
    model = AlexNet()
    img = [[1, 2, 3]] * 10  # 10x3 = 30 < 9216
    codeflash_output = model._extract_features([img]); features = codeflash_output # 1.11μs -> 370ns (201% faster)

def test_extract_features_longer_than_features_size():
    # Test with an image whose flattened size is more than features_size (should truncate)
    model = AlexNet()
    img = [list(range(100)) for _ in range(100)]  # 100x100 = 10000 > 9216
    codeflash_output = model._extract_features([img]); features = codeflash_output # 781ns -> 340ns (130% faster)

def test_extract_features_empty_batch():
    # Test with an empty batch (no images)
    model = AlexNet()
    codeflash_output = model._extract_features([]); features = codeflash_output # 862ns -> 331ns (160% faster)

# --- EDGE TEST CASES ---

def test_extract_features_empty_image():
    # Test with an image that's an empty list
    model = AlexNet()
    img = [[]]
    codeflash_output = model._extract_features(img); features = codeflash_output # 1.17μs -> 321ns (265% faster)




def test_extract_features_image_with_nested_empty_lists():
    # Test with an image containing nested empty lists
    model = AlexNet()
    img = [[], [], []]
    codeflash_output = model._extract_features([img]); features = codeflash_output # 1.34μs -> 430ns (212% faster)

def test_extract_features_image_with_mixed_types():
    # Test with an image containing both lists and ints
    model = AlexNet()
    img = [[1, 2], 3, [4, 5]]
    codeflash_output = model._extract_features([img]); features = codeflash_output # 1.15μs -> 371ns (211% faster)

# --- LARGE SCALE TEST CASES ---

def test_extract_features_large_batch():
    # Test with a large batch of images
    model = AlexNet()
    batch_size = 500  # Large but within 1000
    img = [[i for i in range(36)] for _ in range(256)]
    batch = [img for _ in range(batch_size)]
    codeflash_output = model._extract_features(batch); features = codeflash_output # 7.53μs -> 371ns (1931% faster)
    for f in features:
        pass

def test_extract_features_large_single_image():
    # Test with a single image with a large number of elements
    model = AlexNet()
    img = [[i for i in range(100)] for _ in range(100)]  # 100x100 = 10000 > 9216
    codeflash_output = model._extract_features([img]); features = codeflash_output # 861ns -> 381ns (126% faster)

def test_extract_features_large_batch_varied_sizes():
    # Test with a large batch where each image has different sizes
    model = AlexNet()
    batch = []
    for i in range(100):
        img = [[j for j in range(i+1)] for _ in range(10)]  # 10x(i+1)
        batch.append(img)
    codeflash_output = model._extract_features(batch); features = codeflash_output
    for i, f in enumerate(features):
        expected_flat_len = 10*(i+1)

def test_extract_features_large_sparse_images():
    # Test with large batch of images containing mostly zeros
    model = AlexNet()
    batch_size = 100
    img = [[0]*36 for _ in range(256)]
    batch = [img for _ in range(batch_size)]
    codeflash_output = model._extract_features(batch); features = codeflash_output # 1.69μs -> 401ns (322% faster)
    for f in features:
        pass

def test_extract_features_large_batch_with_empty_images():
    # Test with a large batch where all images are empty
    model = AlexNet()
    batch = [[] for _ in range(500)]
    codeflash_output = model._extract_features(batch); features = codeflash_output # 7.46μs -> 370ns (1917% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.